"""
Agno Human-in-the-Loop (HITL) with FastAPI
完整的用户确认流程示例

安装依赖:
pip install fastapi uvicorn agno openai pydantic
"""

from re import T
import traceback
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import HTMLResponse, StreamingResponse
from fastapi.templating import Jinja2Templates
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from typing import List, Optional, Dict, Any, AsyncGenerator
from datetime import datetime
import uvicorn
import os
import json
import asyncio

# Agno imports
from agno.agent import Agent
from agno.models.openai import OpenAIChat
from agno.tools import tool
from agno.db.postgres import PostgresDb

# 设置 OpenAI API Key (请替换为你的密钥)
os.environ["OPENAI_API_KEY"] = "your-openai-api-key-here"

app = FastAPI(title="Agno HITL Demo")

# 配置Jinja2模板
templates = Jinja2Templates(directory="template")

# ============ 数据模型 ============
class ChatMessage(BaseModel):
    role: str
    content: str

class ChatRequest(BaseModel):
    message: str
    run_id: Optional[str] = None

class ToolConfirmation(BaseModel):
    tool_name: str
    tool_args: Dict[str, Any]
    confirmed: bool

class ChatResponse(BaseModel):
    response: str
    is_paused: bool
    run_id: Optional[str] = None
    tools_requiring_confirmation: Optional[List[Dict[str, Any]]] = None
    timestamp: str

# ============ 定义需要确认的工具 ============

@tool(requires_confirmation=True)
def send_email(recipient: str, subject: str, body: str) -> str:
    print("send_email 工具被调用")
    """发送邮件 - 需要用户确认"""
    return f"邮件已发送给 {recipient}，主题：{subject}"

@tool(requires_confirmation=True)
def delete_file(filename: str) -> str:
    print("delete_file 工具被调用")
    """删除文件 - 需要用户确认"""
    return f"文件 {filename} 已删除"

@tool(requires_confirmation=True)
def make_payment(amount: float, recipient: str) -> str:
    print("make_payment 工具被调用")
    """支付款项 - 需要用户确认"""
    return f"已向 {recipient} 支付 ${amount:.2f}"

@tool()
def search_database(query: str) -> str:
    print("search_database 工具被调用")
    """搜索数据库 - 不需要确认"""
    return f"搜索结果：找到与 '{query}' 相关的 5 条记录"

@tool()
def get_current_time() -> str:
    print("get_current_time 工具被调用")
    """获取当前时间 - 不需要确认"""
    return f"当前时间是 {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}"




db = PostgresDb(
    db_url="postgresql://postgres:123456@localhost:5432/agent",
    memory_table="agent_memory", # Specify the table to store memories
)


# ============ 创建 Agno Agent ============
agent = Agent(
    name="HITL Assistant",
    model=OpenAIChat(
        id="qwen3",
        base_url="http://xai.dyg.com.cn/v1",
        api_key="dyg-y2fyntVOODsN7a1iNUgqjfh6dcBMfpciZWhn4KGJ1x8",
    ),
    tools=[
        send_email,
        delete_file,
        make_payment,
        search_database,
        get_current_time
    ],
    instructions=[
        "你是一个友好的 AI 助手",
        "当用户要求执行敏感操作时，使用相应的工具",
        "对于发送邮件、删除文件、支付等操作，系统会暂停等待用户确认",
        "用中文回复用户"
    ],
    db=db,
    markdown=True
)

# 存储运行状态的字典
run_states: Dict[str, Any] = {}

# ============ API 路由 ============

@app.get("/", response_class=HTMLResponse)
async def read_root(request: Request):
    """返回前端界面"""
    context = {
        "request": request,
        "title": "Agno HITL 演示",
        "app_name": "Agno Human-in-the-Loop 演示",
        "version": "1.0.0",
        "build_time": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        "debug_mode": True,
        "safe_operations": [
            "搜索数据库、获取时间（自动执行）"
        ],
        "dangerous_operations": [
            "发送邮件、删除文件、支付款项（需要确认）"
        ],
        "example_commands": '"发送邮件给 zhang@example.com, 主题:年度会议, 内容:明天下午5点年度会议"，"删除 report.pdf"，"向 John 支付 $100"',
        "welcome_message": "你好！我是 HITL 助手。我可以帮你执行各种操作，对于敏感操作（如发送邮件、删除文件、支付等），我会先征求你的确认。"
    }
    return templates.TemplateResponse("index.html", context)

async def generate_chat_stream(message: str) -> AsyncGenerator[str, None]:
    """生成聊天流式响应"""
    try:
        print(f"Starting chat stream for message: {message}")
        
        # 发送开始事件
        start_data = {
            "type": "start",
            "message": "开始处理请求...",
            "timestamp": datetime.now().isoformat()
        }
        yield f"data: {json.dumps(start_data)}\n\n"
        await asyncio.sleep(0.1)  # 确保前端能接收到开始事件
        
        # 运行 agent
        run_generator = agent.run(message, stream=True)
        print(f"Agent run generator created: {run_generator}")
        
        accumulated_content = ""  # 累积内容，用于计算增量
        has_content = False  # 标记是否已经有内容输出
        
        for run_response in run_generator:
            print(f"Stream event: {run_response.event if hasattr(run_response, 'event') else 'no event'}")
            print(f"Response type: {type(run_response)}")
            print(f"Response attributes: {dir(run_response)}")
            
            # 检查是否需要用户确认
            if hasattr(run_response, 'is_paused') and run_response.is_paused:
                print("Response is paused, extracting tools")
                # 提取需要确认的工具信息
                tools_requiring_confirmation = []
                if hasattr(run_response, 'tools_requiring_confirmation'):
                    for tool in run_response.tools_requiring_confirmation:
                        tools_requiring_confirmation.append({
                            "tool_name": tool.tool_name,
                            "tool_args": tool.tool_args,
                            "confirmed": False
                        })
                
                # 保存运行状态
                run_states[run_response.run_id] = run_response
                
                # 发送暂停事件
                pause_data = {
                    "type": "pause",
                    "run_id": run_response.run_id,
                    "tools_requiring_confirmation": tools_requiring_confirmation,
                    "timestamp": datetime.now().isoformat()
                }
                yield f"data: {json.dumps(pause_data)}\n\n"
                return
            
            # 处理增量内容
            if hasattr(run_response, 'delta') and run_response.delta:
                print(f"Sending delta: {run_response.delta}")
                has_content = True
                delta_data = {
                    "type": "delta",
                    "delta": run_response.delta,
                    "run_id": getattr(run_response, 'run_id', 'unknown'),
                    "timestamp": datetime.now().isoformat()
                }
                yield f"data: {json.dumps(delta_data)}\n\n"
                await asyncio.sleep(0.05)  # 小延迟确保流式效果
            
            # 处理完整内容 - 计算增量
            elif hasattr(run_response, 'content') and run_response.content:
                current_content = run_response.content
                print(f"Current content length: {len(current_content)}, Accumulated length: {len(accumulated_content)}")
                
                if current_content != accumulated_content:
                    # 计算新增的内容
                    if current_content.startswith(accumulated_content):
                        delta = current_content[len(accumulated_content):]
                        if delta:  # 只有当有新内容时才发送
                            print(f"Sending calculated delta: {delta[:50]}...")
                            has_content = True
                            
                            # 如果内容较长，分块发送以实现流式效果
                            if len(delta) > 50:
                                # 将长内容分块发送
                                chunk_size = 10
                                for i in range(0, len(delta), chunk_size):
                                    chunk = delta[i:i+chunk_size]
                                    delta_data = {
                                        "type": "delta",
                                        "delta": chunk,
                                        "run_id": getattr(run_response, 'run_id', 'unknown'),
                                        "timestamp": datetime.now().isoformat()
                                    }
                                    yield f"data: {json.dumps(delta_data)}\n\n"
                                    await asyncio.sleep(0.03)  # 控制流式速度
                            else:
                                delta_data = {
                                    "type": "delta",
                                    "delta": delta,
                                    "run_id": getattr(run_response, 'run_id', 'unknown'),
                                    "timestamp": datetime.now().isoformat()
                                }
                                yield f"data: {json.dumps(delta_data)}\n\n"
                                await asyncio.sleep(0.05)
                            
                            accumulated_content = current_content
                    else:
                        # 如果内容不是累积的，分块发送全部内容
                        print(f"Sending full content as delta: {current_content[:50]}...")
                        has_content = True
                        
                        # 分块发送
                        chunk_size = 15
                        for i in range(0, len(current_content), chunk_size):
                            chunk = current_content[i:i+chunk_size]
                            delta_data = {
                                "type": "delta",
                                "delta": chunk,
                                "run_id": getattr(run_response, 'run_id', 'unknown'),
                                "timestamp": datetime.now().isoformat()
                            }
                            yield f"data: {json.dumps(delta_data)}\n\n"
                            await asyncio.sleep(0.04)  # 控制流式速度
                        
                        accumulated_content = current_content
            
            # 如果没有内容或增量，发送事件信息
            else:
                print(f"No content or delta, sending event info")
                event_data = {
                    "type": "event",
                    "event": str(getattr(run_response, 'event', 'unknown')),
                    "timestamp": datetime.now().isoformat()
                }
                yield f"data: {json.dumps(event_data)}\n\n"
                await asyncio.sleep(0.1)
        
        # 如果没有通过delta或content发送任何内容，尝试从最终响应获取
        if not has_content and hasattr(run_response, 'content') and run_response.content:
            print("No streaming content detected, sending final content as stream")
            content = run_response.content
            chunk_size = 20
            for i in range(0, len(content), chunk_size):
                chunk = content[i:i+chunk_size]
                delta_data = {
                    "type": "delta",
                    "delta": chunk,
                    "run_id": getattr(run_response, 'run_id', 'unknown'),
                    "timestamp": datetime.now().isoformat()
                }
                yield f"data: {json.dumps(delta_data)}\n\n"
                await asyncio.sleep(0.05)
        
        print("Stream completed")
        # 发送完成事件
        complete_data = {
            "type": "complete",
            "timestamp": datetime.now().isoformat()
        }
        yield f"data: {json.dumps(complete_data)}\n\n"
        
    except Exception as e:
        print(f"Error in generate_chat_stream: {e}")
        traceback.print_exc()
        error_data = {
            "type": "error",
            "error": str(e),
            "timestamp": datetime.now().isoformat()
        }
        yield f"data: {json.dumps(error_data)}\n\n"

@app.post("/api/chat/stream")
async def chat_stream(request: ChatRequest):
    """
    处理流式聊天请求
    """
    return StreamingResponse(
        generate_chat_stream(request.message),
        media_type="text/plain",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "Content-Type": "text/event-stream",
        }
    )

@app.post("/api/chat", response_model=ChatResponse)
async def chat(request: ChatRequest):
    """
    处理聊天请求（非流式，保持向后兼容）
    如果工具需要确认，agent 会暂停并返回待确认的工具
    """
    try:
        # 运行 agent
        for run_response in agent.run(request.message, stream=True):
            print(run_response.event, type(run_response.event))
            # 检查是否需要用户确认
            if run_response.is_paused:
                # 提取需要确认的工具信息
                tools_requiring_confirmation = []
                if hasattr(run_response, 'tools_requiring_confirmation'):
                    for tool in run_response.tools_requiring_confirmation:
                        tools_requiring_confirmation.append({
                            "tool_name": tool.tool_name,
                            "tool_args": tool.tool_args,
                            "confirmed": False
                        })
                
                # 保存运行状态
                run_states[run_response.run_id] = run_response
                
                return ChatResponse(
                    response="",
                    is_paused=True,
                    run_id=run_response.run_id,
                    tools_requiring_confirmation=tools_requiring_confirmation,
                    timestamp=datetime.now().isoformat()
                )
            else:
                # 正常响应
                return ChatResponse(
                    response=run_response.content,
                    is_paused=False,
                    run_id=run_response.run_id,
                    timestamp=datetime.now().isoformat()
                )
        
    except Exception as e:
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=str(e))

async def generate_continue_stream(run_id: str, tools: List[Dict]) -> AsyncGenerator[str, None]:
    """生成继续运行的流式响应"""
    try:
        if run_id not in run_states:
            error_data = {
                "type": "error",
                "error": "Run not found",
                "timestamp": datetime.now().isoformat()
            }
            yield f"data: {json.dumps(error_data)}\n\n"
            return
        
        run_response = run_states[run_id]
        
        # 更新工具确认状态
        if hasattr(run_response, 'tools_requiring_confirmation'):
            for i, tool in enumerate(run_response.tools_requiring_confirmation):
                if i < len(tools):
                    tool.confirmed = tools[i].get("confirmed", False)
        print(run_response)
        # 继续执行
        # 根据 Agno 库的 API，continue_run 可能需要不同的参数格式
        try:
            # 方法1: 尝试使用 run_id
            continued_response = agent.continue_run(run_id=run_response.run_id)
        except (AttributeError, TypeError) as e:
            print(f"方法1失败 (run_id): {e}")
            try:
                # 方法2: 尝试直接传递 run_response 对象
                continued_response = agent.continue_run(run_response)
            except (AttributeError, TypeError) as e2:
                print(f"方法2失败 (run_response): {e2}")
                try:
                    # 方法3: 尝试传递 run_response 作为关键字参数
                    continued_response = agent.continue_run(run_response=run_response)
                except (AttributeError, TypeError) as e3:
                    print(f"方法3失败 (run_response=): {e3}")
                    # 方法4: 如果都失败，抛出详细错误信息
                    error_msg = f"无法继续运行。尝试的方法都失败了:\n1. run_id: {e}\n2. run_response: {e2}\n3. run_response=: {e3}"
                    raise Exception(error_msg)
        
        # 检查是否还需要确认
        if continued_response.is_paused:
            tools_requiring_confirmation = []
            if hasattr(continued_response, 'tools_requiring_confirmation'):
                for tool in continued_response.tools_requiring_confirmation:
                    tools_requiring_confirmation.append({
                        "tool_name": tool.tool_name,
                        "tool_args": tool.tool_args,
                        "confirmed": False
                    })
            
            run_states[continued_response.run_id] = continued_response
            
            pause_data = {
                "type": "pause",
                "run_id": continued_response.run_id,
                "tools_requiring_confirmation": tools_requiring_confirmation,
                "timestamp": datetime.now().isoformat()
            }
            yield f"data: {json.dumps(pause_data)}\n\n"
            return
        
        # 发送内容作为增量
        if hasattr(continued_response, 'content') and continued_response.content:
            delta_data = {
                "type": "delta",
                "delta": continued_response.content,
                "run_id": continued_response.run_id,
                "timestamp": datetime.now().isoformat()
            }
            yield f"data: {json.dumps(delta_data)}\n\n"
        
        # 清理状态
        del run_states[run_id]
        
        # 发送完成事件
        complete_data = {
            "type": "complete",
            "timestamp": datetime.now().isoformat()
        }
        yield f"data: {json.dumps(complete_data)}\n\n"
        
    except Exception as e:
        traceback.print_exc()
        error_data = {
            "type": "error",
            "error": str(e),
            "timestamp": datetime.now().isoformat()
        }
        yield f"data: {json.dumps(error_data)}\n\n"

@app.post("/api/continue/stream")
async def continue_run_stream(request: dict):
    """
    继续执行被暂停的运行（流式）
    """
    run_id = request.get("run_id")
    tools = request.get("tools", [])
    
    return StreamingResponse(
        generate_continue_stream(run_id, tools),
        media_type="text/plain",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "Content-Type": "text/event-stream",
        }
    )

@app.post("/api/continue")
async def continue_run(request: dict):
    """
    继续执行被暂停的运行（非流式，保持向后兼容）
    """
    try:
        run_id = request.get("run_id")
        tools = request.get("tools", [])

        print(run_id)
        print(tools)
        
        if run_id not in run_states:
            raise HTTPException(status_code=404, detail="Run not found")
        
        run_response = run_states[run_id]
        
        # 更新工具确认状态
        if hasattr(run_response, 'tools_requiring_confirmation'):
            for i, tool in enumerate(run_response.tools_requiring_confirmation):
                if i < len(tools):
                    tool.confirmed = tools[i].get("confirmed", False)
        
        # 继续执行
        print(run_response.tools)
        print(run_id)
        print(tools)
        
        # 继续执行 - 使用相同的错误处理逻辑
        try:
            # 方法1: 尝试使用 run_id
            continued_response = agent.continue_run(run_id=run_response.run_id)
        except (AttributeError, TypeError) as e:
            print(f"方法1失败 (run_id): {e}")
            try:
                # 方法2: 尝试直接传递 run_response 对象
                continued_response = agent.continue_run(run_response)
            except (AttributeError, TypeError) as e2:
                print(f"方法2失败 (run_response): {e2}")
                try:
                    # 方法3: 尝试传递 run_response 作为关键字参数
                    continued_response = agent.continue_run(run_response=run_response)
                except (AttributeError, TypeError) as e3:
                    print(f"方法3失败 (run_response=): {e3}")
                    # 方法4: 如果都失败，抛出详细错误信息
                    error_msg = f"无法继续运行。尝试的方法都失败了:\n1. run_id: {e}\n2. run_response: {e2}\n3. run_response=: {e3}"
                    raise Exception(error_msg)
        
        # 检查是否还需要确认
        if continued_response.is_paused:
            tools_requiring_confirmation = []
            if hasattr(continued_response, 'tools_requiring_confirmation'):
                for tool in continued_response.tools_requiring_confirmation:
                    tools_requiring_confirmation.append({
                        "tool_name": tool.tool_name,
                        "tool_args": tool.tool_args,
                        "confirmed": False
                    })
            
            run_states[continued_response.run_id] = continued_response
            
            return ChatResponse(
                response="",
                is_paused=True,
                run_id=continued_response.run_id,
                tools_requiring_confirmation=tools_requiring_confirmation,
                timestamp=datetime.now().isoformat()
            )
        
        # 清理状态
        del run_states[run_id]
        
        return ChatResponse(
            response=continued_response.content,
            is_paused=False,
            timestamp=datetime.now().isoformat()
        )
        
    except Exception as e:
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/api/health")
async def health_check():
    """健康检查"""
    return {"status": "healthy", "message": "Agno HITL API is running"}

@app.get("/config", response_class=HTMLResponse)
async def config_page(request: Request):
    """配置页面"""
    context = {
        "request": request,
        "title": "Agno HITL 配置",
        "app_name": "Agno Human-in-the-Loop 配置",
        "version": "1.0.0",
        "build_time": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        "debug_mode": False,  # 配置页面不显示测试按钮
        "safe_operations": [
            "搜索数据库",
            "获取当前时间",
            "查看系统状态"
        ],
        "dangerous_operations": [
            "发送邮件（需要确认）",
            "删除文件（需要确认）",
            "支付款项（需要确认）",
            "修改系统配置（需要确认）"
        ],
        "example_commands": '"配置邮件服务器"，"查看当前配置"，"重置所有设置"',
        "welcome_message": "欢迎来到配置页面！在这里你可以管理系统设置和查看当前配置。所有敏感操作都需要你的确认。"
    }
    return templates.TemplateResponse("index.html", context)

if __name__ == "__main__":
    print("🚀 Starting Agno HITL + FastAPI Server")
    print("📝 请确保已设置 OPENAI_API_KEY 环境变量")
    print("🌐 访问 http://localhost:8000")
    uvicorn.run(app, host="0.0.0.0", port=8000)